import numpy as np
from util.plot import set_axis


# 计算层速度
def calLayerVel(time, depth):
    """
        计算层速度
        :param: time，数组
        :param: depth，数组
        :return: 层速度数组
    """
    stack_velocities = depth / time
    two_way_times = time
    """
    使用Dix公式计算层速度
    :param stack_velocities: 叠加速度列表，单位：米/秒
    :param two_way_times: 双程垂直反射时间列表，单位：秒
    :return: 层速度列表，单位：米/秒
    """
    num_layers = len(stack_velocities)
    layer_velocities = [0,]

    for n in range(1, num_layers):
        vn_squared = (stack_velocities[n] ** 2 * two_way_times[n] -
                      stack_velocities[n - 1] ** 2 * two_way_times[n - 1])
        vn_squared /= (two_way_times[n] - two_way_times[n - 1])
        layer_velocities.append(np.sqrt(vn_squared))

    layer_velocities[0] = layer_velocities[1]

    return np.array(layer_velocities) * 1000  # 将层速度从米/秒转换为毫米/秒


# 绘制层速度
def draw_vel(plt, time, deep):
    """
    绘制层速度
    :param plt: plt
    :param time: 时间
    :param deep: 深度
    :return: ax, fig
    """
    # ============================================================
    # 数据预处理
    # ============================================================

    vel = calLayerVel(np.array(time), np.array(deep))

    # ============================================================
    # 绘制数据
    # ============================================================

    fig = plt.figure(facecolor='w', dpi=300)
    ax = fig.add_axes([0.05, 0.3, 0.35, 0.55])  # 拉平, left，bottom，width，height
    ax.plot(vel, deep, 'green', lw=0.24, label='auto')

    # ============================================================
    # 设置坐标轴
    # ============================================================
    x_min = np.min(vel)
    x_max = np.max(vel)
    y_min = np.min(deep) - 20
    y_max = np.max(deep) + 20

    set_axis(ax, fig, plt, x_min, x_max, y_min, y_max, "Vel(m/ms)", "Deep(m)")

    return ax, fig
